package cn.jbolt.ndemo;

import cn.hutool.core.util.StrUtil;
import cn.jbolt.ai.app.service.AiAppService;
import cn.jbolt.ai.core.chain.ChainContext;
import cn.jbolt.ai.core.chain.node.base.AppConfigConstant;
import cn.jbolt.ai.core.chain.node.base.LlmConfig;
import cn.jbolt.ai.core.chain.node.base.RenderType;
import cn.jbolt.ai.core.chain.node.provider.CommonNodeProvider;
import com.fasterxml.jackson.databind.JsonNode;

import com.jboltai.capability.model.JBoltAI;
import com.jboltai.event.ai.aichat.AIChatEvent;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;


@Component
public class TextExtractorNode implements CommonNodeProvider {

    @Autowired
    private AiAppService aiAppService;

    @Override
    public void provide(JsonNode nodeConfig, ChainContext context) {
        AIChatEvent event = JBoltAI.chat();

        //根据配置初始化event
        configEvent(nodeConfig, context, event);
        event.publish().await();
        if (event.isFail()) {
            context.sendErrorMsg("Ai提取竞品失败:" + event.getError().getMsg());
            throw new RuntimeException(event.getError().getMsg(), event.getError().getException());
        }
        context.sendResponseMsg(event.getResult().getContent(),RenderType.text);
    }

    private LlmConfig configEvent(JsonNode nodeConfig, ChainContext context, AIChatEvent event) {

        LlmConfig config = LlmConfig.from(nodeConfig.get(AppConfigConstant.CONFIG_TYPE_LLM));
        String prompt = generatePrompt(config, context);
        event.prompt(prompt);
        event.setStream(config.isStream());
        if (config.getModelType() != null) {
            event.setModel(config.getModelType());
        }
        if (StrUtil.isNotBlank(config.getDefaultModel())) {
            event.setModelName(config.getDefaultModel());
        }
        if (config.getMaxTokens() > 0) {
            event.setMaxTokens(config.getMaxTokens());
        }
        if (StrUtil.isNotBlank(config.getSystemRole())) {
            event.systemMessage(config.getSystemRole());
        }
        if (config.isUseContext()) {
            //如果需要使用上下文,则要去对应的msg中取出相应的上下文,本demo不需要使用上下文
//            Long appId = context.getAiApp().getId();
//            Long sessionId = context.getAiSession().getId();
//            List<AIMessage> messageList = aiAppService.getAppSessionContextMessages(sessionId, appId);
//            if (messageList != null && messageList.size() > 0) {
//                event.contextMessages(messageList);
//            }
        }
        if (config.getTemperature() != null && config.getTemperature() > 0) {
            event.setTemperature(config.getTemperature());
        }
        if (config.getTopP() != null && config.getTopP() > 0) {
            event.setTopP(config.getTopP());
        }
        if (config.getTopK() != null && config.getTopK() > 0) {
            event.setTopK(config.getTopK());
        }
        return config;

    }

    /**
     * 生成提示词
     *
     * @param config  配置
     * @param context 上下文
     * @return 处理后的提示词
     */
    private String generatePrompt(LlmConfig config, ChainContext context) {

        String template = config.getPromptTpl();
        if (context.getArgs() == null || context.getArgs().isEmpty()) {
            return template;
        }


        for (String key : context.getArgs().keySet()) {
            Object value = context.getArgs().get(key);
            if (value != null) {
                // 处理可能存在的空格情况，如 ${ key } 或 ${key}
                template = template.replaceAll("\\$\\{\\s*" + key + "\\s*\\}", value.toString());
                break;
            }
        }

        return template;
    }
}
